import subprocess
import os
import sys
import re
from typing import List

def get_removable_disks() -> List[str]:
    """获取当前连接的U盘设备列表"""
    try:
        # 执行fdisk命令获取设备信息
        result = subprocess.run(
            ['sudo', 'fdisk', '-l'],
            capture_output=True,
            text=True,
            check=True
        )
        devices = []
        # 精确匹配Disk开头的设备行，并捕获/dev/sdX格式的设备路径
        disk_pattern = re.compile(r'^Disk\s+(/dev/sd[a-z]+):')

        for line in result.stdout.split('\n'):
            line = line.strip()
            if match := disk_pattern.match(line):
                print(line, "---------------")
                device_path = match.group(1)
                print(f"提取到设备: {device_path}")
                # # 验证是否为可移动设备（可选）
                # removable_path = f"/sys/block/{os.path.basename(device_path)}/removable"
                # if os.path.exists(removable_path) and open(removable_path).read().strip() == '1':
                devices.append(device_path)
        print("devices: ", devices)
        return devices
    except subprocess.CalledProcessError as e:
        print(f"\033[31m设备检测失败: {e.stderr}\033[0m")
        return []

def find_valid_partition(device: str) -> str:
    """查找设备的有效分区"""
    try:
        result = subprocess.run(
            ['lsblk', '-o', 'NAME,TYPE', '-n', '-l', device],
            capture_output=True,
            text=True,
            check=True
        )
        # 匹配分区设备（如sdb1）
        for line in result.stdout.splitlines():
            name, type_ = line.strip().split()
            if type_ == 'part':
                return f"/dev/{name}"
        return ""
    except Exception as e:
        print(f"\033[33m分区检测失败: {str(e)}\033[0m")
        return ""

def safe_mount_device(device: str) -> bool:
    """执行安全挂载操作"""
    mount_point = "/mnt/usb"
    try:
        # 清理已有挂载
        subprocess.run(['sudo', 'umount', mount_point], stderr=subprocess.DEVNULL)
        os.makedirs(mount_point, exist_ok=True)

        # 构建挂载命令
        mount_cmd = [
            'sudo', 'mount', '-t', 'auto',
            '-o', 'uid=1000,gid=1000,dmask=022,errors=remount-ro',
            device,
            mount_point
        ]

        # 执行挂载
        subprocess.run(mount_cmd, check=True)
        print(f"\033[32m✔ 成功挂载 {device} 到 {mount_point}\033[0m")
        return True
    except subprocess.CalledProcessError as e:
        error_msg = e.stderr.lower()
        if 'already mounted' in error_msg:
            print(f"\033[33m⚠ 设备已挂载\033[0m")
            return True
        print(f"\033[31m✖ 挂载失败: {e.stderr}\033[0m")
        return False

def handle_usb_storage():
    """处理U盘存储设备"""
    print("\n=== U盘即时检测系统 ===")

    # 获取当前连接的U盘
    disks = get_removable_disks()

    if not disks:
        print("\033[36mℹ 未检测到任何U盘设备\033[0m")
        return

    print(f"发现 {len(disks)} 个U盘设备:")

    for index, disk in enumerate(disks, 1):
        print(f"\n[{index}] 正在检测设备: {disk}")

        # 查找有效分区
        partition = find_valid_partition(disk)
        if partition:
            print(f"  发现分区: {partition}")
            if safe_mount_device(partition):
                return True
                subprocess.Popen(['xdg-open', '/mnt/usb'])
            else:
                return False
        else:
            print("\033[33m⚠ 未找到有效分区\033[0m")
            # 交互式确认是否挂载整个磁盘
            if input("是否尝试挂载整个磁盘？(y/n) ").lower() == 'y':
                if safe_mount_device(disk):
                    return True
                else:
                    return False

    print("\033[31m✖ 所有设备挂载尝试失败\033[0m")

if __name__ == "__main__":
    try:
        # 依赖检查
        if not os.path.exists('/usr/sbin/fdisk'):
            print("请先安装必要工具: sudo apt install fdisk")
            sys.exit(1)

        print(handle_usb_storage())
    except Exception as e:
        print(f"\033[31m✖ 系统错误: {str(e)}\033[0m")
